{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", message=\"numpy.dtype size changed\")\n",
    "warnings.filterwarnings(\"ignore\", message=\"numpy.ufunc size changed\")\n",
    "\n",
    "# basic uses\n",
    "import pprint\n",
    "import re\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle as pk\n",
    "import joblib\n",
    "from itertools import zip_longest\n",
    "from pathlib import Path\n",
    "from math import ceil\n",
    "from random import uniform\n",
    "import seaborn as sb\n",
    "\n",
    "# plotting figures\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Rectangle\n",
    "\n",
    "# RDKit molecule conversion and drawing\n",
    "from rdkit import Chem\n",
    "from rdkit.Chem import Draw\n",
    "from rdkit.Chem.Draw import rdMolDraw2D\n",
    "\n",
    "# modeling\n",
    "import torch\n",
    "from sklearn.metrics import mean_absolute_error\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "# user-friendly print\n",
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = \"all\"\n",
    "\n",
    "sb.set()\n",
    "sb.set_context(\"notebook\", font_scale=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    " # prediction vs. observation plots (single test trial)\n",
    "def draw(y_true, y_pred, y_true_fit=None, y_pred_fit=None, *, prop_name, log_scale=False, file_dir=None, file_name=None):\n",
    "\n",
    "    mask = ~np.isnan(y_pred)\n",
    "    y_true = y_true[mask]\n",
    "    y_pred = y_pred[mask]\n",
    "        \n",
    "    data = pd.DataFrame(dict(Observation=y_true, Prediction=y_pred, dataset=['test'] * len(y_true)))\n",
    "    scores = metrics(data['Observation'], data['Prediction'])\n",
    "    \n",
    "    if y_true_fit is not None and y_pred_fit is not None:\n",
    "        mask = ~np.isnan(y_pred_fit)\n",
    "        y_true_fit = y_true_fit[mask]\n",
    "        y_pred_fit = y_pred_fit[mask]\n",
    "        train_ = pd.DataFrame(dict(Observation=y_true_fit, Prediction=y_pred_fit, dataset=['train'] * len(y_true_fit)))\n",
    "        data = pd.concat([train_, data])\n",
    "\n",
    "    if log_scale:\n",
    "        data = data.apply(lambda c: np.log(c.values) if c.dtype.type is not np.object_ else c, axis=0)\n",
    "#         test_ = test_.apply(lambda c: np.log(c.values) if c.dtype.type is not np.object_ else c, axis=0)\n",
    "\n",
    "\n",
    "\n",
    "#     with sb.set(font_scale=2.5):\n",
    "    g = sb.lmplot(x=\"Prediction\", y=\"Observation\", hue=\"dataset\", ci=None,\n",
    "                  data=data, palette=\"Set1\", height=10, legend=False,  markers=[\".\", \"o\"],\n",
    "                  scatter_kws={'s': 25, 'alpha': 0.7}, hue_order=['train', 'test'])\n",
    "    \n",
    "    ax = plt.gca()\n",
    "    tmp = [data[\"Prediction\"].max(), data[\"Prediction\"].min(), data[\"Observation\"].max(), data[\"Observation\"].max()]\n",
    "    min_, max_ = np.min(tmp), np.max(tmp)\n",
    "    margin = (max_- min_) / 15\n",
    "    min_ = min_ - margin\n",
    "    max_ = max_ + margin\n",
    "    ax.set_xlim(min_, max_)\n",
    "    ax.set_ylim(min_, max_)\n",
    "    ax.set_xlabel(ax.get_xlabel(), fontsize='xx-large')\n",
    "    ax.set_ylabel(ax.get_ylabel(), fontsize='xx-large')\n",
    "    ax.tick_params(axis='both', which='major', labelsize='xx-large')\n",
    "    ax.plot((min_, max_), (min_, max_), ':', color='gray')\n",
    "    ax.set_title(prop_name, fontsize='xx-large')\n",
    "    if log_scale:\n",
    "        ax.set_title(prop_name + ' (log scale)', fontsize='xx-large')\n",
    "    ax.text(0.98, 0.03,\n",
    "            'MAE: %.5f\\nRMSE: %.5f\\nPearsonR: %.5f\\nSpearmanR: %.5f' % (scores['mae'], scores['rmse'], scores['pearsonr'], scores['spearmanr']),\n",
    "            transform=ax.transAxes, horizontalalignment='right', fontsize='xx-large')\n",
    "\n",
    "    ax.legend(loc='upper left', markerscale=2, fancybox=True, shadow=True, frameon=True, facecolor='w', fontsize=18)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    if file_dir and file_name:\n",
    "        if log_scale:\n",
    "            plt.savefig(file_dir + '/' + file_name + '_log_scale.png', dpi=300, bbox_inches='tight')\n",
    "        else:\n",
    "            plt.savefig(file_dir + '/' + file_name + '.png', dip=300, bbox_inches='tight')\n",
    "    else:\n",
    "        print('Missing directory and/or file name information!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prediction vs. observation plots (cross-validation: list of test trials)\n",
    "def draw_cv(y_trues, y_preds, y_trues_fit, y_preds_fit, *, prop_name, log_scale=False, file_dir=None, file_name=None):\n",
    "    cv=len(y_trues)\n",
    "    \n",
    "    y_true = np.concatenate(y_trues)\n",
    "    y_pred = np.concatenate(y_preds)\n",
    "    y_true_fit = np.concatenate(y_trues_fit)\n",
    "    y_pred_fit = np.concatenate(y_preds_fit)\n",
    "    \n",
    "    mask = ~np.isnan(y_pred_fit)\n",
    "    y_true_fit = y_true_fit[mask]\n",
    "    y_pred_fit = y_pred_fit[mask]\n",
    "\n",
    "    mask = ~np.isnan(y_pred)\n",
    "    y_true = y_true[mask]\n",
    "    y_pred = y_pred[mask]\n",
    "        \n",
    "    test_ = pd.DataFrame(dict(Observation=y_true, Prediction=y_pred, dataset=['test'] * len(y_true)))\n",
    "    train_ = pd.DataFrame(dict(Observation=y_true_fit, Prediction=y_pred_fit, dataset=['train'] * len(y_true_fit)))\n",
    "    data = pd.concat([train_, test_])\n",
    "    if log_scale:\n",
    "        data = data.apply(lambda c: np.log(c.values) if c.dtype.type is not np.object_ else c, axis=0)\n",
    "        test_ = test_.apply(lambda c: np.log(c.values) if c.dtype.type is not np.object_ else c, axis=0)\n",
    "\n",
    "    scores = metrics(test_['Observation'], test_['Prediction'])\n",
    "\n",
    "#     sb.set(font_scale=2.5)\n",
    "    _, ax = plt.subplots(figsize=(8, 8), dpi=150)\n",
    "    ax = sb.lmplot(x=\"Prediction\", y=\"Observation\", hue=\"dataset\", ci=None,\n",
    "                  data=data, palette=\"Set1\", height=10, legend=False,  markers=[\".\", \"o\"],\n",
    "                  scatter_kws={'s': 25, 'alpha': 0.7}, hue_order=['train', 'test'], ax=ax)\n",
    "    \n",
    "#     ax = plt.gca()\n",
    "    tmp = [data[\"Prediction\"].max(), data[\"Prediction\"].min(), data[\"Observation\"].max(), data[\"Observation\"].max()]\n",
    "    min_, max_ = np.min(tmp), np.max(tmp)\n",
    "    margin = (max_- min_) / 15\n",
    "    min_ = min_ - margin\n",
    "    max_ = max_ + margin\n",
    "    ax.set_xlim(min_, max_)\n",
    "    ax.set_ylim(min_, max_)\n",
    "    ax.set_xlabel(ax.get_xlabel(), fontsize='xx-large')\n",
    "    ax.set_ylabel(ax.get_ylabel(), fontsize='xx-large')\n",
    "    ax.tick_params(axis='both', which='major', labelsize='xx-large')\n",
    "    ax.plot((min_, max_), (min_, max_), ':', color='gray')\n",
    "    ax.set_title(prop_name, fontsize='xx-large')\n",
    "    if log_scale:\n",
    "        ax.set_title(prop_name + ' (log scale)', fontsize='xx-large')\n",
    "    ax.text(0.98, 0.03,\n",
    "            '$%d-fold$ CV\\nmae: %.5f\\nrmse: %.5f\\npearsonr: %.5f\\nspearmanr: %.5f' % (cv, scores['mae'], scores['rmse'], scores['pearsonr'], scores['spearmanr']),\n",
    "            transform=ax.transAxes, horizontalalignment='right', fontsize='xx-large')\n",
    "    ax.legend(loc='upper left', markerscale=2, fancybox=True, shadow=True, frameon=True, facecolor='w')\n",
    "    plt.tight_layout()\n",
    "    if file_dir and file_name:\n",
    "        if log_scale:\n",
    "            plt.savefig(file_dir + '/' + file_name + '_log_scale.png', dpi=300, bbox_inches='tight')\n",
    "        else:\n",
    "            plt.savefig(file_dir + '/' + file_name + '.png', dip=300, bbox_inches='tight')\n",
    "    else:\n",
    "        print('Missing directory and/or file name information!')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculating basic statistics for predictions\n",
    "def metrics(y_true, y_pred, ignore_nan=True):\n",
    "    from sklearn.metrics import mean_absolute_error, r2_score, mean_squared_error\n",
    "    from scipy.stats import pearsonr, spearmanr\n",
    "    \n",
    "    if ignore_nan:\n",
    "        mask = ~np.isnan(y_pred)\n",
    "        y_true = y_true[mask]\n",
    "        y_pred = y_pred[mask]\n",
    "    \n",
    "    mae = mean_absolute_error(y_true, y_pred)\n",
    "    rmse = np.sqrt(mean_squared_error(y_true, y_pred))\n",
    "    r2 = r2_score(y_true, y_pred)\n",
    "    pr, p_val = pearsonr(y_true, y_pred)\n",
    "    sr, _ = spearmanr(y_true, y_pred)\n",
    "    return dict(\n",
    "        mae=mae,\n",
    "        rmse=rmse,\n",
    "        r2=r2,\n",
    "        pearsonr=pr,\n",
    "        spearmanr=sr,\n",
    "        p_value=p_val\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "def quicksort(x):\n",
    "    if x==[]: return []\n",
    "\n",
    "    smallerSorted = quicksort([a for a in x[1:] if os.path.getmtime(str(a)) <= os.path.getmtime(str(x[0]))])\n",
    "    biggerSorted = quicksort([a for a in x[1:] if os.path.getmtime(str(a)) > os.path.getmtime(str(x[0]))])\n",
    "\n",
    "    return(smallerSorted+[x[0]]+biggerSorted)\n",
    "\n",
    "def retrieve_nn_models_from_with_cv(*props, score='pearsonr', cv=10):\n",
    "    for prop in props:\n",
    "        p = Path(prop)\n",
    "        models = [x for x in p.iterdir() if x.is_dir() and x.name != '.ipynb_checkpoints']\n",
    "        models = quicksort(models)\n",
    "        return [Checker.load(x.name, x.parent) for x in models[:10]], p"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
